#-------------------------------------------------------------------------------
# Name:        module1
# Purpose:     引入边界条件到[K]中,并返回解[U],[P]
#               input:
#                   K:全局刚度矩阵,(M,M) numpy.array
#                   BcDict:位移约束,key (int) = 自由度序号(1-based) , value (float) = 自由度约束值
#                   LoadDict:节点力加载,key (int) = 自由度序号(1-based) , value (float) = 施加的节点力加载或者等效节点力加载
#
# Author:      Administrator
#
# Created:     08-03-2025
# Copyright:   (c) Administrator 2025
# Licence:     <your licence>
#-------------------------------------------------------------------------------
import numpy as np
from typing import Dict,List,Tuple
import scipy as sc
def Method1(K:np.ndarray,BcDict:Dict[int,float],LoadDict:Dict[int,float])->Tuple:
    """划行划列法"""
    dofNum=K.shape[0]
    # 初始化向量
    U,P=np.zeros((dofNum,1)),np.zeros((dofNum,1))
    prescribedDofIndexs=np.array(list(BcDict.keys()))-1
    #使用集合运算,全部自由度与约束自由度求差, 得到自由位移自由度的
    freeDofIndexs=np.array(list(set(range(dofNum))-set(prescribedDofIndexs.tolist())),dtype=int)
    # 已知节点力加到P
    for label,Pval in LoadDict.items():
        ind=label-1
        P[ind,0]+=Pval
    # 已知节点位移(prescribed dof)
    for label,Uval in BcDict.items():
        ind=label-1
        U[ind,0]+=Uval
    U2=U[np.ix_(prescribedDofIndexs,[0])].copy()
    # 已知节点力(free dof)
    P1=P[np.ix_(freeDofIndexs,[0])].copy()
    # 重新划分K行列
    K11=K[np.ix_(freeDofIndexs,freeDofIndexs)].copy()
    K12=K[np.ix_(freeDofIndexs,prescribedDofIndexs)].copy()
    K21=K[np.ix_(prescribedDofIndexs,freeDofIndexs)].copy()
    K22=K[np.ix_(prescribedDofIndexs,prescribedDofIndexs)].copy()
    # 计算自由节点位移值
    U1=np.dot(sc.linalg.inv(K11),P1-K12.dot(U2))
    # 计算支反力
    P2=np.dot(K21,U1)+np.dot(K22,U2)
    # 合并到U,P向量
    U[np.ix_(freeDofIndexs,[0])]=U1
    P[np.ix_(prescribedDofIndexs,[0])]=P2
    return U,P

def Method2(K:np.ndarray,BcDict:Dict[int,float],LoadDict:Dict[int,float])->Tuple:
    K_origin=K.copy()
    dofNum=K.shape[0]
    # 初始化向量
    U,P=np.zeros((dofNum,1)),np.zeros((dofNum,1))
    # 已知节点力加到 P
    for label,Pval in LoadDict.items():
        ind=label-1
        P[ind,0]+=Pval
    # 循环所有的位移约束
    for label,Uval in BcDict.items():
        j=label-1
        #Step1
        for i in range(dofNum):
            P[i,0]=P[i,0]-K[i,j]*Uval
        #Step2
        for i in range(dofNum):
            K[i,j]=0
            K[j,i]=0
        K[j,j]=1
        #Step3
        P[j,0]=Uval
    # 求解 K'U=P'
    U_=sc.linalg.solve(K,P)
    P_=np.dot(K_origin,U_)
    return U_,P_

def Method3(K:np.ndarray,BcDict:Dict[int,float],LoadDict:Dict[int,float])->Tuple:
    C=np.max(K)*10e6
    K_origin=K.copy()
    dofNum=K.shape[0]
    # 初始化向量
    U,P=np.zeros((dofNum,1)),np.zeros((dofNum,1))
    # 已知节点力加到 P
    for label,Pval in LoadDict.items():
        ind=label-1
        P[ind,0]+=Pval
    # 循环所有位移约束
    for label,Uval in BcDict.items():
        j=label-1
        # Step1
        K[j,j]=K[j,j]*C
        # Step2
        P[j,0]=K[j,j]*Uval
    # 求解 K'U=P'
    U_=sc.linalg.solve(K,P)
    P_=np.dot(K_origin,U_)
    return U_,P_

def Method4(K:np.ndarray,BcDict:Dict[int,float],LoadDict:Dict[int,float])->Tuple:
    if np.any(np.array(list(BcDict.values())) != 0):
        raise ValueError('该方法不能处理非0位移加载')
    K_origin=K.copy()
    dofNum=K.shape[0]
    # 初始化向量
    U,P=np.zeros((dofNum,1)),np.zeros((dofNum,1))
    # 已知节点力加到 P
    for label,Pval in LoadDict.items():
        ind=label-1
        P[ind,0]+=Pval
    # loop all nodal bcs
    for label, Uval in BcDict.items():
        j=label-1
        K[j,:]=0.0
        K[:,j]=0.0
        K[j,j]=1.0
        P[j,0]=0
    # solve K'U=P'
    U_=sc.linalg.solve(K,P)
    P_=np.dot(K_origin,U_)
    return U_,P_

if __name__ == '__main__':
    K=np.array([[12,600,-12,600],
                [600,40000,-600,20000],
                [-12,-600,12,-600],
                [600,20000,-600,40000]])
    Bcs={1:0,2:0}
    loads={3:-50,4:20}

    # 精确解
    extract_U=np.array([0,0,-16.5667,-0.2480])
    extract_P=np.array([50.0,4980.0,-50,20])
    # 求解
    u,p=Method4(K,Bcs,loads)

##    print(f"extract U=\n{extract_U}")
##    print(f"extract_P=\n{extract_P}")
    print('solving by method 4')
    print(f"u=\n{u.T}")
    print(f"p=\n{p.T}")
